#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@163.com
@desc: 
"""
import argparse
from pathlib import Path

import cv2
import numpy as np
from dltools.utils.io import read_voc_xml
from dltools.utils.log import get_console_logger
from object_detection.utils import visualization_utils as vis_utils

from utils.file_read import GoodsReader


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--image-root', type=str, default='', dest='root',
        help='The directory where the image data and the annotation is stored.'
             '保存图像与标记数据的根目录, 图像与其相应标记应当在同一级目录。')
    parser.add_argument(
        '--output-root', type=str, default='', dest='output',
        help='the root for all the output.'
             '输出结果的根目录')
    return parser.parse_args()


def _create_label_list():
    label_list = []
    for idx in range(100):
        label_list.append('package%d' % idx)
    return label_list


def _create_category_index():
    category_index = {}
    for idx in range(100):
        category_index[idx + 1] = {'id': idx + 1, 'name': 'package%d' % idx}
    return category_index


def read_xml_as_box_label(xml, shape):
    xml, _ = read_voc_xml(xml, shape, _create_label_list())
    labels = []
    boxes = []
    cls_name = []
    for ob in xml['objects']:
        labels.append(ob['label'])
        boxes.append(ob['box'])
        cls_name.append(ob['name'])
    labels = np.asarray(labels) + 1
    boxes = np.asarray(boxes)
    return boxes, labels, cls_name


def draw_objects_in_image(data_reader, output):
    for data in data_reader:
        logger.info('Draw image {}'.format(data['image']))
        image = cv2.imread(data['image'])
        boxes, labels, cls_name = read_xml_as_box_label(
            data['xml'], image.shape)
        image = vis_utils.visualize_boxes_and_labels_on_image_array(
            image, boxes, labels, np.ones_like(labels, dtype=np.float16),
            _create_category_index(), max_boxes_to_draw=100
        )
        output_name = Path(output) / Path(data['image']).name
        cv2.imwrite(str(output_name), image)


if __name__ == '__main__':
    ARGS = parse_args()
    logger = get_console_logger('Draw Goods')
    _data_reader = GoodsReader(ARGS.root, logger)
    draw_objects_in_image(_data_reader, ARGS.output)
